from typing import Dict
import ccxt.pro as ccxtpro
import asyncio
import time
import logging

class PriceMonitor:
    def __init__(self, proxy: str = None):
        # 初始化日志
        logging.basicConfig(
            level=logging.INFO,
            format='%(asctime)s - %(levelname)s - %(message)s'
        )
        self.logger = logging.getLogger(__name__)
        
        # 价格缓存
        self.spot_prices: Dict[str, Dict] = {}
        self.futures_prices: Dict[str, Dict] = {}
        
        # 修改价格缓存结构
        self.orderbook_data = {
            # 'Exchange': {
            #     'spot': {
            #         'symbol': {
            #             'bids': [[price, amount], ...],
            #             'asks': [[price, amount], ...],
            #             'timestamp': timestamp
            #         }
            #     },
            #     'futures': {
            #         'symbol': {
            #             'bids': [[price, amount], ...],
            #             'asks': [[price, amount], ...],
            #             'timestamp': timestamp
            #         }
            #     }
            # }
        }
        
        # 保存配置为类属性，以便后续重用
        self.base_config = {
            'enableRateLimit': True,
            'timeout': 30000,
            'options': {
                'defaultType': 'spot',
                'watchTicker': {
                    'name': 'ticker'
                }
            },
            'ws': {
                'options': {
                    'keepAlive': True,
                    'pingInterval': 15000,
                    'reconnectTimeout': 5000,
                    'reconnectDelay': 5000,
                }
            }
        }
        
        # 如果有代理，添加代理配置
        if proxy:
            self.base_config['proxyUrl'] = proxy
            
        # 初始化所有交易所
        self.exchanges = {}
        
        # Bitget
        self.exchanges['Bitget'] = {
            'spot': ccxtpro.bitget({
                **self.base_config,
                'options': {'defaultType': 'spot'}
            }),
            'futures': ccxtpro.bitget({
                **self.base_config,
                'options': {'defaultType': 'swap'}
            })
        }
        
        # OKX
        self.exchanges['OKX'] = {
            'spot': ccxtpro.okx({
                **self.base_config,
                'options': {'defaultType': 'spot'}
            }),
            'futures': ccxtpro.okx({
                **self.base_config,
                'options': {'defaultType': 'swap'}
            })
        }
        
        # Binance
        self.exchanges['Binance'] = {
            'spot': ccxtpro.binance({
                **self.base_config,
                'options': {'defaultType': 'spot'}
            }),
            'futures': ccxtpro.binance({
                **self.base_config,
                'options': {'defaultType': 'future'}
            })
        }
        
        # Gate
        self.exchanges['Gate'] = {
            'spot': ccxtpro.gateio({
                **self.base_config,
                'options': {'defaultType': 'spot'}
            }),
            'futures': ccxtpro.gateio({
                **self.base_config,
                'options': {'defaultType': 'swap'}
            })
        }

        # 初始化每个交易所的数据结构
        for exchange_name in ['Bitget', 'OKX', 'Binance', 'Gate']:
            self.orderbook_data[exchange_name] = {
                'spot': {},
                'futures': {}
            }

    async def initialize(self):
        """初始化交易所连接和市场数据"""
        max_retries = 3
        retry_delay = 5
        
        for attempt in range(max_retries):
            try:
                self.logger.info(f"正在加载市场数据... (尝试 {attempt + 1}/{max_retries})")
                
                # 分别加载市场数据，避免一起失败
                for exchange_name, exchange_clients in self.exchanges.items():
                    for client_type, client in exchange_clients.items():
                        try:
                            await client.load_markets()
                            self.logger.info(f"{exchange_name} {client_type} 市场数据加载完成")
                        except Exception as e:
                            self.logger.error(f"{exchange_name} {client_type} 市场数据加载失败: {str(e)}")
                
                return
                
            except Exception as e:
                self.logger.error(f"初始化失败 (尝试 {attempt + 1}/{max_retries}): {str(e)}")
                if attempt < max_retries - 1:
                    await asyncio.sleep(retry_delay)
                else:
                    raise e

    def convert_symbol_to_futures(self, symbol: str, exchange: str) -> str:
        """转换现货符号为期货符号"""
        base = symbol.split('/')[0]
        
        if exchange == 'Bitget':
            return symbol
        elif exchange == 'OKX':
            return f"{base}-USDT-SWAP"
        elif exchange == 'Binance':
            return f"{base}USDT"
        elif exchange == 'Gate':
            return f"{base}_USDT"
        return symbol

    async def watch_exchange_depth(self, exchange_name: str, symbols: list):
        """监控单个交易所的深度数据"""
        exchange = self.exchanges[exchange_name]
        
        try:
            tasks = []
            # 同时订阅所有交易对
            for symbol in symbols:
                # 订阅现货深度
                self.logger.info(f"订阅 {exchange_name} 现货深度 {symbol}")
                spot_task = asyncio.create_task(self.process_orderbook(
                    exchange['spot'], 
                    symbol, 
                    f"{exchange_name}_{symbol}", 
                    is_spot=True
                ))
                tasks.append(spot_task)
                
                # 订阅期货深度
                futures_symbol = self.convert_symbol_to_futures(symbol, exchange_name)
                self.logger.info(f"订阅 {exchange_name} 期货深度 {futures_symbol}")
                futures_task = asyncio.create_task(self.process_orderbook(
                    exchange['futures'], 
                    futures_symbol, 
                    f"{exchange_name}_{symbol}", 
                    is_spot=False
                ))
                tasks.append(futures_task)
            
            # 等待所有订阅任务完成
            await asyncio.gather(*tasks)
                
        except Exception as e:
            self.logger.error(f"{exchange_name} 深度订阅错误: {str(e)}")
            self.logger.exception(e)

    async def process_orderbook(self, exchange, symbol: str, key: str, is_spot: bool):
        """处理订单簿数据"""
        exchange_name = key.split('_')[0]
        market_type = 'spot' if is_spot else 'futures'
        
        while True:
            try:
                while True:
                    orderbook = await exchange.watch_order_book(symbol)
                    if orderbook and orderbook['asks'] and orderbook['bids']:
                        # 存储完整深度数据
                        self.orderbook_data[exchange_name][market_type][symbol] = {
                            'bids': orderbook['bids'],
                            'asks': orderbook['asks'],
                            'timestamp': orderbook['timestamp']
                        }
                    
            except Exception as e:
                self.logger.error(f"处理深度数据错误 {symbol}: {str(e)}")
                self.logger.exception(e)
                await asyncio.sleep(5)

    def print_price_diff(self):
        """打印价格差异"""
        print("\n" + "="*100)
        print(f"时间: {time.strftime('%Y-%m-%d %H:%M:%S')}")
        print("-"*100)
        print(f"{'交易所':<15} {'交易对':<15} {'现货买一':<15} {'现货卖一':<15} "
              f"{'期货买一':<15} {'期货卖一':<15} {'价差(%)':<15}")
        print("-"*100)
        
        for exchange_name in self.orderbook_data:
            spot_data = self.orderbook_data[exchange_name]['spot']
            futures_data = self.orderbook_data[exchange_name]['futures']
            
            for symbol in spot_data:
                spot_orderbook = spot_data[symbol]
                # 获取对应的期货symbol
                futures_symbol = self.convert_symbol_to_futures(symbol, exchange_name)
                futures_orderbook = futures_data.get(futures_symbol)
                
                if spot_orderbook and futures_orderbook:
                    # 获取买一卖一价格
                    spot_bid = float(spot_orderbook['bids'][0][0])
                    spot_ask = float(spot_orderbook['asks'][0][0])
                    futures_bid = float(futures_orderbook['bids'][0][0])
                    futures_ask = float(futures_orderbook['asks'][0][0])
                    
                    # 计算中间价
                    spot_price = (spot_bid + spot_ask) / 2
                    futures_price = (futures_bid + futures_ask) / 2
                    
                    # 计算价差
                    price_diff = (futures_price - spot_price) / spot_price * 100
                    
                    # 显示原始symbol（去掉USDT后缀等）
                    display_symbol = symbol.split('/')[0]
                    
                    print(f"{exchange_name:<15} {display_symbol:<15} "
                          f"{spot_bid:<15.8f} {spot_ask:<15.8f} "
                          f"{futures_bid:<15.8f} {futures_ask:<15.8f} "
                          f"{price_diff:>15.4f}")
        
        print("="*100 + "\n")

    async def run(self, symbols: list):
        """运行监控"""
        try:
            self.logger.info("正在启动价格监控...")
            self.logger.info(f"监控交易对: {symbols}")
            
            # 初始化
            await self.initialize()
            
            # 启动所有交易所的价格订阅任务
            tasks = []
            for exchange_name in self.exchanges:
                tasks.append(self.watch_exchange_depth(exchange_name, symbols))
            
            # 启动价格打印任务
            async def print_task():
                while True:
                    self.print_price_diff()
                    await asyncio.sleep(3)
            
            tasks.append(print_task())
            
            # 运行所有任务
            await asyncio.gather(*tasks)
            
        except KeyboardInterrupt:
            self.logger.info("\n正在退出程序...")
        except Exception as e:
            self.logger.error(f"运行错误: {str(e)}")
            self.logger.exception(e)
        finally:
            # 关闭所有交易所连接
            for exchange in self.exchanges.values():
                for client in exchange.values():
                    await client.close()

    async def get_funding_rates(self, exchange_name: str = None, symbols: list = None):
        """获取交易所的资金费率
        
        Args:
            exchange_name: 交易所名称，如果为None则获取所有交易所
            symbols: 交易对列表，如果为None则获取所有交易对
            
        Returns:
            按资金费率降序排序的资金费率字典
        """
        all_funding_rates = {}
        
        try:
            # 确定要查询的交易所列表
            exchanges_to_query = [exchange_name] if exchange_name else self.exchanges.keys()
            
            for exch_name in exchanges_to_query:
                if exch_name not in self.exchanges:
                    self.logger.error(f"交易所 {exch_name} 不存在")
                    continue
                    
                exchange = self.exchanges[exch_name]['futures']
                
                try:
                    # OKX 特殊处理
                    if exch_name == 'OKX':
                        self.logger.info(f"使用特殊方法获取 OKX 资金费率...")
                        # 先获取所有可交易的合约
                        markets = await exchange.load_markets()
                        swap_symbols = [symbol for symbol in markets.keys() 
                                      if markets[symbol]['swap'] and 'USDT' in symbol]
                        
                        for symbol in swap_symbols:
                            try:
                                # 单独获取每个合约的资金费率
                                funding_rate_data = await exchange.fetch_funding_rate(symbol)
                                
                                if not isinstance(funding_rate_data, dict):
                                    continue
                                    
                                rate = funding_rate_data.get('fundingRate', 0)
                                
                                # 跳过零费率
                                if rate == 0:
                                    continue
                                    
                                # 提取基础币种
                                if '/' in symbol:
                                    base_symbol = symbol.split('/')[0]
                                elif '-' in symbol:
                                    base_symbol = symbol.split('-')[0]
                                else:
                                    base_symbol = symbol.replace('USDT', '').replace('SWAP', '')
                                
                                # 标准化交易对格式
                                std_symbol = f"{base_symbol}/USDT"
                                
                                # 如果指定了交易对列表，只保留列表中的交易对
                                if symbols and std_symbol not in symbols:
                                    continue
                                
                                all_funding_rates[std_symbol] = {
                                    'rate': float(rate),
                                    'timestamp': int(time.time() * 1000),
                                    'next_funding_time': funding_rate_data.get('nextFundingTime', 0),
                                    'exchange': exch_name
                                }
                                
                                self.logger.info(f"获取 {exch_name} {symbol} 资金费率: {rate}")
                            except Exception as e:
                                self.logger.error(f"获取 {exch_name} {symbol} 资金费率失败: {str(e)}")
                    else:
                        # 其他交易所使用批量获取
                        self.logger.info(f"正在获取 {exch_name} 所有资金费率...")
                        all_rates = await exchange.fetch_funding_rates()
                        self.logger.info(f"成功获取 {exch_name} 资金费率，共 {len(all_rates)} 个")
                        
                        # 处理获取到的资金费率
                        for futures_symbol, rate_info in all_rates.items():
                            # 检查 rate_info 是否为字典
                            if not isinstance(rate_info, dict):
                                self.logger.warning(f"{exch_name} {futures_symbol} 资金费率数据格式异常: {rate_info}")
                                continue
                                
                            # 过滤非USDT交易对
                            if not ('USDT' in futures_symbol or 'usdt' in futures_symbol):
                                continue
                                
                            # 提取基础币种
                            if '/' in futures_symbol:
                                base_symbol = futures_symbol.split('/')[0]
                            elif '-' in futures_symbol:
                                base_symbol = futures_symbol.split('-')[0]
                            elif '_' in futures_symbol:
                                base_symbol = futures_symbol.split('_')[0]
                            else:
                                base_symbol = futures_symbol.replace('USDT', '').replace('usdt', '')
                            
                            # 标准化交易对格式
                            std_symbol = f"{base_symbol}/USDT"
                            
                            # 如果指定了交易对列表，只保留列表中的交易对
                            if symbols and std_symbol not in symbols:
                                continue
                            
                            # 获取资金费率 - 修复可能的数据结构问题
                            if 'fundingRate' in rate_info:
                                rate = rate_info['fundingRate']
                            elif 'rate' in rate_info:
                                rate = rate_info['rate']
                            elif isinstance(rate_info, (int, float)):
                                rate = rate_info
                            else:
                                self.logger.warning(f"无法从 {rate_info} 中提取资金费率")
                                rate = 0
                            
                            # 跳过零费率
                            if rate == 0.0:
                                continue

                            # 获取下次结算时间
                            if 'nextFundingTime' in rate_info:
                                next_funding_time = rate_info['nextFundingTime']
                            elif 'nextFundingTimestamp' in rate_info:
                                next_funding_time = rate_info['nextFundingTimestamp']
                            else:
                                next_funding_time = 0
                            
                            # 存储资金费率
                            all_funding_rates[std_symbol] = {
                                'rate': float(rate),
                                'timestamp': int(time.time() * 1000),
                                'next_funding_time': next_funding_time,
                                'exchange': exch_name
                            }
                            
                            self.logger.info(f"获取 {exch_name} {futures_symbol} 资金费率: {rate}")
                            
                            # 更新深度数据中的资金费率
                            if futures_symbol in self.orderbook_data[exch_name]['futures']:
                                self.orderbook_data[exch_name]['futures'][futures_symbol]['funding_rate'] = float(rate)
                    
                except Exception as e:
                    self.logger.error(f"获取 {exch_name} 所有资金费率失败: {str(e)}")
                    
                    # 如果批量获取失败，尝试单个获取（作为备选方案）
                    if symbols:
                        self.logger.info(f"尝试单个获取 {exch_name} 资金费率...")
                        for symbol in symbols:
                            try:
                                futures_symbol = self.convert_symbol_to_futures(symbol, exch_name)
                                funding_rate_data = await exchange.fetch_funding_rate(futures_symbol)
                                
                                # 检查返回的数据结构
                                if not isinstance(funding_rate_data, dict):
                                    self.logger.warning(f"{exch_name} {futures_symbol} 单个资金费率数据格式异常: {funding_rate_data}")
                                    continue
                                
                                # 获取资金费率
                                if 'fundingRate' in funding_rate_data:
                                    rate = funding_rate_data['fundingRate']
                                elif 'rate' in funding_rate_data:
                                    rate = funding_rate_data['rate']
                                else:
                                    self.logger.warning(f"无法从 {funding_rate_data} 中提取资金费率")
                                    continue
                                
                                # 跳过零费率
                                if rate == 0.0:
                                    continue
                                
                                all_funding_rates[symbol] = {
                                    'rate': float(rate),
                                    'timestamp': int(time.time() * 1000),
                                    'next_funding_time': funding_rate_data.get('nextFundingTime', 0),
                                    'exchange': exch_name
                                }
                                
                                self.logger.info(f"获取 {exch_name} {futures_symbol} 资金费率: {rate}")
                                
                                # 更新深度数据中的资金费率
                                if futures_symbol in self.orderbook_data[exch_name]['futures']:
                                    self.orderbook_data[exch_name]['futures'][futures_symbol]['funding_rate'] = float(rate)
                                
                            except Exception as e:
                                self.logger.error(f"获取 {exch_name} {futures_symbol} 资金费率失败: {str(e)}")
            
            # 按资金费率降序排序
            sorted_funding_rates = dict(sorted(
                all_funding_rates.items(), 
                key=lambda item: item[1]['rate'], 
                reverse=True
            ))
            
            return sorted_funding_rates
            
        except Exception as e:
            self.logger.error(f"获取资金费率失败: {str(e)}")
            return {}

async def test_funding_rates():
    """测试获取资金费率功能"""
    print("\n" + "="*100)
    print("测试获取资金费率功能")
    print("="*100)
    
    # 创建监控实例 - 不使用代理
    monitor = PriceMonitor(proxy=None)
    
    try:
        # 初始化交易所连接
        await monitor.initialize()
        
        # 测试获取特定交易所特定币种的资金费率
        # print("\n1. 获取 Binance 的 BTC/USDT 资金费率:")
        # rates = await monitor.get_funding_rates("Binance", ["BTC/USDT"])
        # print_funding_rates(rates)
        
        # 测试获取特定交易所所有币种的资金费率
        # print("\n2. 获取 OKX 的所有币种资金费率 (仅显示前5个):")
        # rates = await monitor.get_funding_rates("OKX")
        # print_funding_rates(rates, limit=5)
        
        # 测试获取所有交易所特定币种的资金费率
        print("\n3. 获取所有交易所的 ETH/USDT 资金费率:")
        rates = await monitor.get_funding_rates(symbols=["ETH/USDT"])
        print_funding_rates(rates)
        
        # 测试获取所有交易所所有币种的资金费率
        print("\n4. 获取所有交易所所有币种资金费率 (仅显示前10个):")
        rates = await monitor.get_funding_rates()
        print_funding_rates(rates, limit=10)
        
    except Exception as e:
        print(f"测试资金费率功能出错: {str(e)}")
        import traceback
        print(traceback.format_exc())
    finally:
        # 关闭所有交易所连接
        for exchange in monitor.exchanges.values():
            for client in exchange.values():
                await client.close()

def print_funding_rates(rates, limit=None):
    """打印资金费率"""
    if not rates:
        print("未获取到资金费率数据")
        return
    
    # 确保资金费率已按降序排序
    sorted_rates = dict(sorted(
        rates.items(), 
        key=lambda item: item[1]['rate'], 
        reverse=True
    ))
        
    print(f"{'交易对':<15} {'交易所':<10} {'资金费率(%)':<15} {'下次结算时间':<20}")
    print("-"*60)
    
    count = 0
    for symbol, data in sorted_rates.items():
        # 格式化下次结算时间
        next_time = data.get('next_funding_time', 0)
        if next_time:
            next_time_str = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(next_time/1000))
        else:
            next_time_str = "未知"
            
        # 显示资金费率（百分比），保留四位小数
        rate_percent = data['rate'] * 100
        
        # 根据资金费率的正负值添加符号，使用固定格式确保不出现科学计数法
        if rate_percent > 0:
            rate_str = f"+{rate_percent:.4f}%"
        elif rate_percent < 0:
            rate_str = f"{rate_percent:.4f}%"
        else:
            rate_str = f"{rate_percent:.4f}%"
        
        print(f"{symbol:<15} {data['exchange']:<10} {rate_str:>15} {next_time_str:<20}")
        
        count += 1
        if limit and count >= limit:
            print(f"\n... 仅显示前 {limit} 条记录，共 {len(sorted_rates)} 条 ...")
            break
    
    print("-"*60)

async def main():
    # 要监控的交易对
    symbols = ['BTC/USDT', 'ETH/USDT','SOL/USDT']
    
    # 测试资金费率功能
    await test_funding_rates()
    
    # 创建监控实例 - 不使用代理
    # monitor = PriceMonitor(proxy=None)
    #
    # try:
    #     # 运行监控
    #     await monitor.run(symbols)
    # except Exception as e:
    #     print(f"运行错误: {e}")
    #     import traceback
    #     print(traceback.format_exc())

if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        print("\n程序已停止")
    except Exception as e:
        print(f"程序出错: {str(e)}") 